import os
import subprocess
from multiprocessing import Process, Event
import signal
import atexit
import sys
import time

def run_server(stop_event, server_name, host, port, cpu, ram, rom):
    def signal_handler(signum, frame):
        stop_event.set()

    # 注册信号处理函数
    signal.signal(signal.SIGTERM, signal_handler)

    # 设置环境变量
    env = os.environ.copy()
    env['VIM_SERVER_NAME'] = server_name
    env['VIM_RESOURCE_CPU'] = cpu
    env['VIM_RESOURCE_RAM'] = ram
    env['VIM_RESOURCE_ROM'] = rom

    # 构建uvicorn命令
    command = ['uvicorn','run:app',
        '--host', host,
        '--port', port,
        '--log-level', 'error',
    ]

    with subprocess.Popen(command, env=env) as proc:
        try:
            while not stop_event.is_set():
                time.sleep(0.5)
        except Exception as e:
            print(f"Exception in server process: {e}")
        finally:
            proc.terminate()
            proc.wait()  # 等待子进程结束

def run_clean():
    command = ['python','clean.py','-db']
    with subprocess.Popen(command) as proc:
        proc.wait()

def create_vim_process(stop_event, **kwargs):
    vimName = kwargs.get('vimName','vim_default')
    vimIp = kwargs.get('vimIp','127.0.0.1')
    vimPort = kwargs.get('vimPort','5000')
    vimCpu = kwargs.get('vimCpu','4')
    vimRam = kwargs.get('vimRam','1')
    vimRom = kwargs.get('vimRom','32')
    vimProcess = Process(target=run_server, args=(stop_event, vimName, vimIp, vimPort, 
                                                  vimCpu, vimRam, vimRom))
    vimProcess.daemon = True
    vimProcess.start()

    return vimProcess

def create_clean_process(vim_processes):
    for p in vim_processes:
        p.terminate()
        p.join()
    run_clean()

if __name__ == '__main__':
    stop_event1 = Event()
    stop_event2 = Event()

    vim1Info = {'vimName':'vim1','vimIp':'192.168.31.79','vimPort':'5000',
                'vimCpu':'4','vimRam':'1','vimRom':'128'}
    vim2Info = {'vimName':'vim2','vimIp':'192.168.31.79','vimPort':'5001',
                'vimCpu':'4','vimRam':'1','vimRom':'128'}

    vim1Process = create_vim_process(stop_event1, **vim1Info)
    vim2Process = create_vim_process(stop_event2, **vim2Info)
    vim_processes = [vim1Process, vim2Process]

    # 注册退出处理器
    atexit.register(create_clean_process, vim_processes=vim_processes)

    try:
        while True:
            print('Main process is running...')
            time.sleep(2)
    except KeyboardInterrupt:
        print('Main process received interrupt, initiating clean up...')
        stop_event1.set()
        stop_event2.set()
        create_clean_process(vim_processes)
        print('Cleanup completed, exiting...')
        sys.exit(0)

